﻿using FlyIL;
using FlyIL.Emu;
using System;
using System.Collections.Generic;
using System.Linq;


using static FlyIL.Opt;

namespace FuckVmp.VMP
{


    internal enum HandleType
    {

        none = 0,
        _byte1 = 1,
        _byte2 = 2,
        _byte4 = 4,
        _byte8 = 8,

        nop = 0x10,
        push_vsp = 0x11,

        _push_imm_x = 0x20,
        push_imm1 = 0x21,
        push_imm2 = 0x22,
        push_imm4 = 0x24,
        push_imm8 = 0x28,


        _push_tmpx = 0x30,
        push_tmp1 = 0x31,
        push_tmp2 = 0x32,
        push_tmp4 = 0x34,
        push_tmp8 = 0x38,

        _pop_tmpx = 0x40,
        pop_tmp1 = 0x41,
        pop_tmp2 = 0x42,
        pop_tmp4 = 0x44,
        pop_tmp8 = 0x48,


        _stack_addx = 0x50,
        stack_add1 = 0x51,
        stack_add2 = 0x52,
        stack_add4 = 0x54,
        stack_add8 = 0x58,

        _stack_xchgx = 0x60,
        stack_xchg1 = 0x62,
        stack_xchg2 = 0x62,
        stack_xchg4 = 0x64,
        stack_xchg8 = 0x68,

        _readx = 0x70,
        read1 = 0x71,
        read2 = 0x72,
        read4 = 0x74,
        read8 = 0x78,


        _norx = 0x80,
        not_or1 = 0x81,
        not_or2 = 0x82,
        not_or4 = 0x84,
        not_or8 = 0x88,

        _nandx = 0x90,
        not_and1 = 0x91,
        not_and2 = 0x92,
        not_and4 = 0x94,
        not_and8 = 0x98,


        _shr = 0x100,
        shr1 = 0x101,
        shr2 = 0x102,
        shr4 = 0x104,
        shr8 = 0x108,


        _push_sp = 0x110,
        push_sp4 = 0x114,
        push_sp8 = 0x118,


        rdtsc,
        cpuid,


        ret,


    }

    public class HandleInfo
    {


        /// <summary>
        /// 指令集
        /// </summary>
        public Context Context;

        public List<Inst> insts;


        /// <summary>
        /// 关键指令  
        /// </summary>
        public Inst KeyInst;

        internal HandleType Type;

        public Action<Cpu> action;

        public ulong moniterValue;

        public HandleInfo(Context vcontext)
        {
            Context = vcontext;
            insts = Context.Codes.ToList();
        }


    }

    public class HandleDetecor
    {



        #region 数据 

        byte Arch;

        /// <summary> vmcode 所在  </summary>
        public RegCode RegVip { get; private set; }

        /// <summary> 虚拟栈 所在  </summary>
        public RegCode RegVsp { get; private set; }


        public HandleDetecor(byte arch, RegCode vip, RegCode vsp)
        {
            this.Arch = arch;
            this.RegVip = vip;
            this.RegVsp = vsp;
        }



        #endregion







        /// <summary>
        /// 检测单栈头操作 
        /// </summary>
        bool CheckType1(HandleInfo info, Inst AckTmpIndex, Inst readTop, Inst read, Inst write, HandleType len)
        {
            Inst key;

            if (AckTmpIndex != null)
            {
                if (read != null)
                {
                    ///③ 保存 某个寄存器X // mov qword ptr ss:[rsp+r9],rcx
                    //key = info.Codes.Find(o => o.Mnemonic == Mnemonic.Mov
                    //    && o.Op1Register != Register.None
                    //    && o.MemoryBase == Register.RSP
                    //    && o.MemoryIndex != Register.None);


                    key = info.insts.Find(o => o.Opt == mov && o.O2.IsRegister && o.O1.IsStackValue && o.O1.Index != RegCode.None);

                    if (key != null)
                    {
                        info.KeyInst = key;
                        info.Type = (HandleType._pop_tmpx | len);

                        info.action = (Cpu) =>
                        {
                            if (info.KeyInst.O1.Index.BitSize() == 64)
                                info.moniterValue = Cpu.GetReg64(info.KeyInst.O1.Index);
                            else if (info.KeyInst.O1.Index.BitSize() == 32)
                                info.moniterValue = Cpu.GetReg32(info.KeyInst.O1.Index);

                        };

                        return true;
                    }
                }

                if (write != null)
                {
                    ///③ 读取 某个寄存器X         // mov rbp,[rsp+rdi]
                    key = info.insts.Find(o => o.Opt == mov
                        && o.O2.IsStackValue
                        && o.O2.Index != RegCode.None);
                    if (key != null)
                    {
                        info.KeyInst = key;
                        info.Type = HandleType._push_tmpx | len;


                        info.action = (Cpu) =>
                        {
                            if (this.Arch == 64)
                                info.moniterValue = Cpu.GetReg64(info.KeyInst.O2.Index);
                            else if (this.Arch == 32)
                                info.moniterValue = Cpu.GetReg32(info.KeyInst.O2.Index);

                        };

                        return true;
                    }
                }

            }




            /// 单纯的写入 
            if (read == null)
            {

                var transferVsp = info.insts.Find(o => o.Opt == mov && o.O2.IsReg(RegVsp));

                /// 读取 esp , 并写入栈头 
                if (transferVsp != null && transferVsp.O1 == write.O2)
                {

                    info.Type = this.Arch == 64 ? HandleType.push_sp8 : HandleType.push_sp4;

                    return true;

                }


                ///① 从 vmcode 中读取 8个字节       // mov rbp,qword ptr ds:[rbx] 
                if (info.insts.Find(o => o.Opt == mov && o.O2.Base == RegVip) != null)
                {
                    info.KeyInst = write;
                    info.Type = HandleType._push_imm_x | len;

                    info.action = (Cpu) =>
                    {
                        if (info.KeyInst.O2.BitSize == 64)
                            info.moniterValue = Cpu.GetReg64(info.KeyInst.O2.Base);
                        else if (info.KeyInst.O2.BitSize == 32)
                            info.moniterValue = Cpu.GetReg32(info.KeyInst.O2.Base);
                        else if (info.KeyInst.O2.BitSize == 16)
                            info.moniterValue = Cpu.GetReg16(info.KeyInst.O2.Base);
                        else if (info.KeyInst.O2.BitSize == 8)
                            info.moniterValue = Cpu.GetReg8(info.KeyInst.O2.Base);

                    };

                    return true;
                }

            }


            /// 即读取了栈头 , 又写入的栈头 8 
            if (readTop != null && write != null)
            {
                RegCode right = readTop.O1.Base;
                RegCode left = write.O2.Base;


                var inst = info.insts.Find(o => o.Opt == mov && o.O2.Base == right && o.O2.IsPtr && o.O1.Base == left);

                if (inst != null)
                {
                    info.Type = HandleType._readx | len;
                    return true;
                }

            }


            /// 判断 VRet 
            if (readTop != null)
            {

                var inst = info.insts.Find(o => o.Opt == mov && o.O2 == readTop.O1 && o.O1.Base == RegVip);

                if (inst != null)
                {
                    info.Type = HandleType.ret;
                    return true;
                }
            }


            return false;
        }


        /// <summary>
        /// 检测 栈计算  
        /// </summary>
        bool CheckType2(HandleInfo info, Inst read1, Inst read2, Inst write2, HandleType len)
        {

            var reg = read1.O1;

            // 目前看来，先取出来的 放左值   
            var key = info.insts.Find(o => o.OpsCount == 2 && o.O1 == read1.O1 && o.O2 == read2.O1);


            if (key.Opt == add)
            {
                info.Type = HandleType._stack_addx | len;
                return true;
            }
            else if (key.Opt == or)
            {
                info.Type = HandleType._norx | len;
                return true;
            }
            else if (key.Opt == and)
            {
                info.Type = HandleType._nandx | len;
                return true;
            }
            else if (key.Opt == shr)
            {
                info.Type = HandleType._shr | len;
                return true;
            }
            //info.insts.

            return false;
        }



        /// <summary>
        /// 判断 是 哪个 vm 指令， 并标记关键指令
        /// </summary>
        public bool CheckType(HandleInfo info)
        {

            /// 确认 虚拟栈 编号 
            Inst AckTmpIndex;

            /// 栈头读写 8字节 
            Inst write_1st8 = null, read_1st8 = null;

            /// 栈头读写 第2个8字节 
            Inst read_2nd8 = null, write_2nd8 = null;

            /// byte [vsp + 8 ] 
            Inst read_2nd8_1 = null;

            /// 栈头读写 4字节 
            Inst write_1st4, read_1st4;

            /// byte [vsp + 4 ] 
            Inst read_2nd4_1 = null;

            /// 栈头读写 第2个 4字节 
            Inst write_2nd4 = null, read_2nd4 = null;

            /// 栈头读写 2字节 
            Inst write_1st2, read_1st2;

            /// byte [vsp + 4 ] 
            Inst read_2nd2_1 = null;

            /// 栈头读写 第2个 2字节 
            Inst write_2nd2 = null, read_2nd2 = null;


            /// 栈头读写 1字节 
            Inst write_1st1, read_1st1;

            /// 栈头读写 第2个 1字节 
            Inst write_2nd1 = null, read_2nd1 = null;


            Inst transferVip, transferVsp;

            transferVip = info.insts.Find(o => o.Opt == mov && o.O2.IsReg(RegVip));

            transferVsp = info.insts.Find(o => o.Opt == mov && o.O2.IsReg(RegVsp));


            ///  取计临时栈 X  , 确认是 tmp操作   // movzx r9d,byte ptr ds:[rbx] 
            AckTmpIndex = info.insts.Find(o => o.Opt == movzx && o.O1.BitSize == 32 && o.O2.IsPurePtr(RegVip, 1, 0));


            if (this.Arch == 64)
            {
                // 向栈头写入 8位   mov qword ptr ds:[r11],rbp  // 不是基于 rsp , 无内存偏移
                read_1st8 = info.insts.Find(o => o.Opt == mov && o.O1.IsRegister && o.O2.IsPurePtr(RegVsp, 8, 0));
                write_1st8 = info.insts.Find(o => o.Opt == mov && o.O2.IsRegister && o.O1.IsPurePtr(RegVsp, 8, 0));

                if (read_1st8 != null)
                {
                    read_2nd8_1 = info.insts.Find(o => o.Opt == mov && o.O2.IsPurePtr(RegVsp, 1, 8));
                    if (read_2nd8_1 == null)
                        // 从栈头+8 读取     //mov rdx,qword ptr ds:[r11+8]  // 不是基于 rsp , 带偏移
                        read_2nd8 = info.insts.Find(o => o.Opt == mov && o.O2.IsPurePtr(RegVsp, 8, 8));
                    if (read_2nd8_1 != null || read_2nd8 != null)
                        write_2nd8 = info.insts.Find(o => o.Opt == mov && o.O1.IsPurePtr(RegVsp, 8, 8));
                }
            }

            //  mov edx,[r8] 
            read_1st4 = info.insts.Find(o => o.Opt == mov && o.O1.IsRegister && o.O2.IsPurePtr(RegVsp, 4, 0));
            write_1st4 = info.insts.Find(o => o.Opt == mov && o.O1.IsPurePtr(RegVsp, 4, 0));


            if (read_1st4 != null)
            {
                read_2nd4_1 = info.insts.Find(o => o.Opt == mov && o.O2.IsPurePtr(RegVsp, 1, 4));
                if (read_2nd4_1 == null)
                    // 从栈头+8 读取     //mov rdx,qword ptr ds:[r11+8]  // 不是基于 rsp , 带偏移
                    read_2nd4 = info.insts.Find(o => o.Opt == mov && o.O2.IsPurePtr(RegVsp, 4, 4));
                if (read_2nd4_1 != null || read_2nd4 != null)
                    write_2nd4 = info.insts.Find(o => o.Opt == mov && o.O1.IsPurePtr(RegVsp, 4, 4));
            }


            read_1st2 = info.insts.Find(o => o.Opt == mov && o.O1.IsRegister && o.O2.IsPurePtr(RegVsp, 2, 0));
            write_1st2 = info.insts.Find(o => o.Opt == mov && o.O1.IsPurePtr(RegVsp, 2, 0));

            if (read_1st2 != null)
            {
                read_2nd2_1 = info.insts.Find(o => o.Opt == mov && o.O2.IsPurePtr(RegVsp, 1, 2));
                if (read_2nd4_1 == null)
                    read_2nd2 = info.insts.Find(o => o.Opt == mov && o.O1.IsRegister && o.O2.IsPurePtr(RegVsp, 2, 2));
                if (read_2nd2_1 != null || read_2nd2 != null)
                    write_2nd2 = info.insts.Find(o => o.Opt == mov && o.O1.IsPtr && o.O1.IsPurePtr(RegVsp, 2, 2));
            }


            read_1st1 = info.insts.Find(o => o.Opt == mov && o.O1.IsRegister && o.O2.IsPurePtr(RegVsp, 1, 0));
            write_1st1 = info.insts.Find(o => o.Opt == mov && o.O1.IsPurePtr(RegVsp, 1, 0));
            if (read_1st1 != null)
            {
                read_2nd1 = info.insts.Find(o => o.Opt == mov && o.O1.IsRegister && o.O2.IsPurePtr(RegVsp, 1, 4));
                if (read_2nd1 != null)
                    write_2nd1 = info.insts.Find(o => o.Opt == mov && o.O1.IsPurePtr(RegVsp, 1, 4));
            }


            /**********************
             * 优先判断 emit 指令 
            **********************/

            var emit = info.insts.Find(o => o.Opt == Opt.emit);
            if (emit != null)
            {
                ;

                if (emit.Emit == "Cpuid" && info.insts.Find(o => o.Opt == sub && o.O1.Base == RegVsp && o.O2.u32 == (Arch == 64 ? 0x20 : 0xC)) != null)
                {
                    info.Type = HandleType.cpuid;
                    return true;
                }


            }


            /// 读取栈头 
            Inst readTop = this.Arch == 64 ? read_1st8 : read_1st4;


            /**********************
             * 优先判断  小字节的操作 
            **********************/
            if (read_2nd1 != null || write_2nd1 != null)
                return CheckType2(info, read_1st4, read_2nd1, write_2nd1, HandleType._byte1);
            else if (read_1st1 != null || write_1st1 != null)
                return CheckType1(info, AckTmpIndex, readTop, read_1st1, write_1st1, HandleType._byte1);


            else if (read_2nd2 != null || write_2nd2 != null)
                return CheckType2(info, read_1st2, read_2nd2 != null ? read_2nd2 : read_2nd2_1, write_2nd2, HandleType._byte2);
            else if (read_1st2 != null || write_1st2 != null)
                return CheckType1(info, AckTmpIndex, readTop, read_1st2, write_1st2, HandleType._byte2);


            else if (read_2nd4 != null || write_2nd4 != null)
                return CheckType2(info, read_1st4, read_2nd4 != null ? read_2nd4 : read_2nd4_1, write_2nd4, HandleType._byte4);
            else if (read_1st4 != null || write_1st4 != null)
                return CheckType1(info, AckTmpIndex, readTop, read_1st4, write_1st4, HandleType._byte4);

            else if (read_2nd8 != null || write_2nd8 != null)
                return CheckType2(info, read_1st8, read_2nd8 != null ? read_2nd8 : read_2nd8_1, write_2nd8, HandleType._byte8);
            else if (read_1st8 != null || write_1st8 != null)
                return CheckType1(info, AckTmpIndex, readTop, read_1st8, write_1st8, HandleType._byte8);








            info.Type = HandleType.none;
            return false;


        }




    }
}
